﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using ManagedCuda;
using ManagedCuda.VectorTypes;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Devices;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Data;

namespace RayDen.RayEngine.IntersectionEngines.Cuda
{
    public class CudaFilm : WorkingQueueDevice<CudaSampleBuffer>, Core.Interface.IImageFilm, IDisposable
    {
        private readonly ThreadLocal<object> splat_lock = new ThreadLocal<object>(()=>new object());

        public const int GammaTableSize = 1024;
        public const int FilterTableSize = 16;
        public const int SampleBufferSize = 65536;

        FrameBuffer frameBuffer;
        public int Width, Height;
        public int Texture;

        protected ulong samplesProcessed;
        protected NVContext Context;

        private CudaDeviceVariable<float> gammaBuffer;
        private CudaDeviceVariable<SamplePixel> sampleFrameBuff;
        private CudaDeviceVariable<RgbSpectrum> frameBuff;

        protected List<CudaSampleBuffer> sampleBuffers;
        protected CudaKernel clearFBKernel, clearSampleFBKernel, addSampleBufferKernel, updateFrameBufferKernel;


        public CudaFilm(int w, int h)
        {
            sampleBuffers = new List<CudaSampleBuffer>();
            Width = w;
            Height = h;
            this.Context = new NVContext() { Context = new CudaContext(CudaContext.GetMaxGflopsDeviceId()) };
            this.Init();
        }


        public void ResetProcessor()
        {
            throw new System.NotImplementedException();
        }

        
        public SampleBuffer GetFreeSampleBuffer()
        {
            lock (splat_lock)
            {
                if (this.doneBuffer.Count > 3)
                {
                    CudaSampleBuffer buffer;
                    this.doneBuffer.TryDequeue(out buffer);
                    return buffer;
                }
                else
                {
                    for (var i = 0; i < 5; i++)
                    {
                        var osb = new CudaSampleBuffer(SampleBufferSize);
                        doneBuffer.Enqueue(osb);
                        this.sampleBuffers.Add(osb);
                    }
                    var sb = new CudaSampleBuffer(SampleBufferSize);
                    sampleBuffers.Add(sb);
                    return sb;
                }
            }
        }

        public void FreeSampleBuffer(SampleBuffer buffer)
        {
            lock (splat_lock.Value)
            {
                doneBuffer.Enqueue((CudaSampleBuffer)buffer);
            }
        }


        public void SplatSampleBuffer(bool isPreviewOver, SampleBuffer sampleBuffer)
        {
            lock (splat_lock.Value)
            {
                Context.Context.SetCurrent();
                using (var sb = new CudaSampleBuffer(sampleBuffer.Size))
                {
                    sb.Write(sampleBuffer.Samples);
                    addSampleBufferKernel.Run(Width, Height, sampleFrameBuff.DevicePointer, sampleBuffer.Size,
                                              sb.DeviceBuffer.DevicePointer);
                }
            }

            this.doneBuffer.Enqueue((CudaSampleBuffer)sampleBuffer);
        }

        protected override void Process(CudaSampleBuffer data)
        {
            throw new System.NotImplementedException();
        }

        public void Consume(PathSamplerBase path)
        {
            throw new System.NotImplementedException();
        }

        private RgbSpectrum[] data;
        public void UpdateScreenBuffer()
        {
            lock (splat_lock.Value)
            {
                //PixelUpdateFrameBuffer(const unsigned int width,const unsigned int height,SamplePixel *sampleFrameBuffer,Pixel* frameBuffer,float* gammaTable)
                this.Context.Context.SetCurrent();
                this.updateFrameBufferKernel.Run((uint)Width, (uint)Height, this.sampleFrameBuff.DevicePointer, this.frameBuff.DevicePointer, gammaBuffer.DevicePointer);
                
                if (data == null)
                    data = new RgbSpectrum[Width*Height];
                this.frameBuff.CopyToHost(data);
                if (frameBuffer == null)
                    this.frameBuffer = new FrameBuffer(Width, Height, data); 
                else 
                    frameBuffer.FromRgb(data);
            }
        }

        public void Reset()
        {
            throw new System.NotImplementedException();
        }

        public float[] GetScreenBuffer()
        {
            return this.frameBuffer.GetPixelData();
        }

        private void Init()
        {
            using (var sr = new StreamReader(@"G:\Git\RayDen\CudaMegaRay\x64\Release\cudafilm.cu.ptx"))
            {
                this.clearFBKernel = Context.LoadKernel(sr.BaseStream, "PixelClearFB");
                this.clearFBKernel.BlockDimensions = new dim3(8, 8, 1);
                this.clearFBKernel.GridDimensions = new dim3(Width / 8, Height / 8, 1);
            }

            using (var sr = new StreamReader(@"G:\Git\RayDen\CudaMegaRay\x64\Release\cudafilm.cu.ptx"))
            {
                this.clearSampleFBKernel = Context.LoadKernel(sr.BaseStream, "PixelClearSampleFB");
                //this.clearFBKernel.G = 64;
                this.clearSampleFBKernel.BlockDimensions = new dim3(8, 8, 1);
                this.clearSampleFBKernel.GridDimensions = new dim3(Width / 8, Height / 8, 1);
            }

            using (var sr = new StreamReader(@"G:\Git\RayDen\CudaMegaRay\x64\Release\cudafilm.cu.ptx"))
            {
                this.updateFrameBufferKernel = Context.LoadKernel(sr.BaseStream, "PixelUpdateFrameBuffer");
                this.updateFrameBufferKernel.BlockDimensions = new dim3(8, 8, 1);
                this.updateFrameBufferKernel.GridDimensions = new dim3(Width / 8, Height / 8, 1);
            }

            using (var sr = new StreamReader(@"G:\Git\RayDen\CudaMegaRay\x64\Release\cudafilm.cu.ptx"))
            {
                this.addSampleBufferKernel = Context.LoadKernel(sr.BaseStream, "PixelAddSampleBuffer");
                this.addSampleBufferKernel.BlockDimensions = new dim3(64, 1, 1);
                this.addSampleBufferKernel.GridDimensions = new dim3(Width, 1, 1);
            }
            this.sampleFrameBuff = new CudaDeviceVariable<SamplePixel>(Width*Height);
            this.frameBuff = new CudaDeviceVariable<RgbSpectrum>(Width * Height);

            float x = 0f;
            float gamma = RgbSpectrum.GammaValue;
            float dx = 1f / GammaTableSize;
            float[] gammaTable = new float[GammaTableSize];
            for (var i = 0; i < GammaTableSize; ++i, x += dx)
                gammaTable[i] = (float)Math.Pow(MathLab.Clamp(x, 0f, 1f), 1f / gamma);
            gammaBuffer = gammaTable;

            this.ClearFrameBuffer();
            this.ClearSampleBuffer();
        }

        private void ClearSampleBuffer()
        {
            clearSampleFBKernel.Run(Width, Height, sampleFrameBuff.DevicePointer);
        }

        private void ClearFrameBuffer()
        {
            clearFBKernel.Run(Width, Height, frameBuff.DevicePointer);
        }

        public string GetStats()
        {
            return string.Empty;
        }

        public void SavePNG(string fn, string[] watermark)
        {
            throw new NotImplementedException();
        }

        public void SaveRaw(string fn)
        {
            throw new NotImplementedException();
        }

        int IImageFilm.Width
        {
            get { return Width; }
        }

        int IImageFilm.Height
        {
            get { return this.Height; }
        }

        public void Dispose()
        {
            if (this.Context != null)
            {
                this.Context.Dispose();
            }
            this.gammaBuffer.Dispose();            
            this.frameBuff.Dispose();
            this.sampleFrameBuff.Dispose();
            foreach (var cudaSampleBuffer in sampleBuffers)
            {
                cudaSampleBuffer.Dispose();
            }
        }
    }
}